import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt


def low_pass_filter(n1, img_path):
    # 读取图像
    img = cv.imread(img_path, 0)

    # 傅里叶变换
    dft = cv.dft(np.float32(img), flags=cv.DFT_COMPLEX_OUTPUT)
    fshift = np.fft.fftshift(dft)

    # 设置低通滤波器
    rows, cols = img.shape
    crow, ccol = int(rows / 2), int(cols / 2)  # 中心位置
    mask = np.zeros((rows, cols, 2), np.uint8)
    mask[crow - n1:crow + n1, ccol - n1:ccol + n1] = 1

    # 掩膜图像和频谱图像乘积
    f = fshift * mask

    # 傅里叶逆变换
    ishift = np.fft.ifftshift(f)
    iimg = cv.idft(ishift)
    res = cv.magnitude(iimg[:, :, 0], iimg[:, :, 1])

    # 将结果归一化并转换为uint8类型
    res = cv.normalize(res, None, 0, 255, cv.NORM_MINMAX)
    res = np.uint8(res)

    return res


def high_pass_filter(n, img_path):
    # 读取图像
    img = cv.imread(img_path, 0)

    if img is None:
        print("Error: Image not found or unable to load.")
        return None

    # 傅里叶变换
    f = np.fft.fft2(img)
    fshift = np.fft.fftshift(f)

    # 设置高通滤波器
    rows, cols = img.shape
    crow, ccol = int(rows / 2), int(cols / 2)

    # 应用高通滤波器
    fshift[crow - n:crow + n, ccol - n:ccol + n] = 0

    # 傅里叶逆变换
    ishift = np.fft.ifftshift(fshift)
    iimg = np.fft.ifft2(ishift)
    iimg = np.abs(iimg)

    # 将结果归一化并转换为uint8类型
    iimg = cv.normalize(iimg, None, 0, 255, cv.NORM_MINMAX)
    iimg = np.uint8(iimg)

    return iimg


def apply_fourier_transform(img_path):
    # 读取图像
    img = cv.imread(img_path, 0)

    # 傅里叶变换
    dft = cv.dft(np.float32(img), flags=cv.DFT_COMPLEX_OUTPUT)
    dftshift = np.fft.fftshift(dft)
    res1 = 20 * np.log(cv.magnitude(dftshift[:, :, 0], dftshift[:, :, 1]))

    # 傅里叶逆变换
    ishift = np.fft.ifftshift(dftshift)
    iimg = cv.idft(ishift)
    res2 = cv.magnitude(iimg[:, :, 0], iimg[:, :, 1])

    # 将结果归一化并转换为uint8类型
    res1 = cv.normalize(res1, None, 0, 255, cv.NORM_MINMAX)
    res1 = np.uint8(res1)
    res2 = cv.normalize(res2, None, 0, 255, cv.NORM_MINMAX)
    res2 = np.uint8(res2)

    return res1, res2
#
# # 傅里叶调用函数并传递参数
# if __name__ == "__main__":
#     img_path = 'lena.png'
#
#     n2 = int(input("Enter the size of the filter (n2): "))
#     fourier_img, inverse_fourier_img = apply_fourier_transform(n2, img_path)
#
#     if fourier_img is not None and inverse_fourier_img is not None:
#         print("Fourier transform applied.")
#         print("Fourier image dtype:", fourier_img.dtype)
#         print("Inverse Fourier image dtype:", inverse_fourier_img.dtype)
#
#     # 显示结果图像
#     original_img = cv.imread(img_path, 0)
#
#     plt.subplot(131), plt.imshow(original_img, 'gray'), plt.title('Original Image')
#     plt.axis('off')
#     if fourier_img is not None:
#         plt.subplot(132), plt.imshow(fourier_img, 'gray'), plt.title('Fourier Image')
#         plt.axis('off')
#     if inverse_fourier_img is not None:
#         plt.subplot(133), plt.imshow(inverse_fourier_img, 'gray'), plt.title('Inverse Fourier Image')
#         plt.axis('off')
#
#     plt.show()
# -----------------------高通滤波与低通滤波---------------------------
# img_path = "lena.png"
# low = low_pass_filter(20, img_path)
# high = high_pass_filter(20, img_path)
# print(low.dtype)
# print(high.dtype)
# # 调用函数并传递参数
# if __name__ == "__main__":
#     img_path = 'lena.png'
#
#     n1 = int(input("Enter the size of the low-pass filter (n1): "))
#     low_pass_img = low_pass_filter(n1, img_path)
#     if low_pass_img is not None:
#         print("Low pass filter applied.")
#         print(low_pass_img.dtype)
#
#     n = int(input("Enter the size of the high-pass filter (n): "))
#     high_pass_img = high_pass_filter(n, img_path)
#     if high_pass_img is not None:
#         print("High pass filter applied.")
#         print(high_pass_img.dtype)
#
#     # 显示结果图像
#     original_img = cv.imread(img_path, 0)
#
#     cv.imshow('Original Image', original_img)
#     if low_pass_img is not None:
#         cv.imshow('Low Pass Filtered Image', low_pass_img)
#     if high_pass_img is not None:
#         cv.imshow('High Pass Filtered Image', high_pass_img)
#
#     # 等待按键操作，关闭所有窗口
#     cv.waitKey(0)
#     cv.destroyAllWindows()
